from pymilvus import (
    connections,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
    db,
    utility
)
import numpy as np

# 配置
milvus_alias = "similarity_demo"
db_name = "test_db"
collection_name = "similarity_collection"
dim = 64  # 向量维度


class MilvusVectorManager:
    """Milvus向量管理器，提供完整的增删改查功能"""

    def __init__(self):
        self.setup_milvus()

    def setup_milvus(self):
        """初始化Milvus连接和环境"""
        try:
            # 建立连接
            connections.connect(
                alias=milvus_alias,
                host='127.0.0.1',
                port='19530'
            )
            print("Milvus连接成功")

            # 检查并清理旧环境
            existing_dbs = db.list_database(using=milvus_alias)
            if db_name in existing_dbs:
                db.using_database(db_name, using=milvus_alias)
                # 删除该数据库中的所有集合
                collections = utility.list_collections(using=milvus_alias)
                for coll_name in collections:
                    try:
                        collection = Collection(coll_name, using=milvus_alias)
                        collection.drop()
                        print(f"已删除集合: {coll_name}")
                    except Exception as e:
                        print(f"删除集合 {coll_name} 失败: {e}")
                # 删除数据库
                db.drop_database(db_name, using=milvus_alias)
                print(f"已删除数据库: {db_name}")

            # 创建新数据库
            db.create_database(db_name, using=milvus_alias)
            db.using_database(db_name, using=milvus_alias)
            print(f"已创建并切换到数据库: {db_name}")

        except Exception as e:
            print(f"Milvus设置失败: {e}")
            raise

    def create_collection(self, metric_type, collection_name=None):
        """创建指定相似度度量的集合"""
        try:
            # 定义字段
            fields = [
                FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
                FieldSchema(name="label", dtype=DataType.INT64),  # 用于验证匹配效果的标签
                FieldSchema(name="desc", dtype=DataType.VARCHAR, max_length=100)  # 向量描述
            ]

            # 创建集合
            collection_full_name = collection_name or f"{collection_name}_{metric_type.lower()}"
            schema = CollectionSchema(fields, description=f"使用{metric_type}的集合")
            collection = Collection(
                name=collection_full_name,
                schema=schema,
                using=milvus_alias
            )

            # 创建索引
            index_params = {
                "index_type": "IVF_FLAT",
                "metric_type": metric_type,
                "params": {"nlist": 100}
            }
            collection.create_index("vector", index_params)
            print(f"集合 {collection_full_name} 创建成功")
            return collection
        except Exception as e:
            print(f"创建集合失败: {e}")
            raise

    def insert_data(self, collection, vectors, labels=None, descs=None):
        """插入数据"""
        try:
            # 准备插入数据
            if labels is None:
                labels = [0] * len(vectors)
            if descs is None:
                descs = [""] * len(vectors)

            insert_data = [vectors, labels, descs]
            collection.insert(insert_data)
            collection.flush()
            collection.load()
            print(f"已插入 {len(vectors)} 条数据")
            return True
        except Exception as e:
            print(f"插入数据失败: {e}")
            raise

    def delete_data(self, collection, ids):
        """删除指定ID的数据"""
        try:
            # 根据ID删除数据
            expr = f"id in {ids}"
            collection.delete(expr)
            collection.flush()
            print(f"已删除 {len(ids)} 条数据")
            return True
        except Exception as e:
            print(f"删除数据失败: {e}")
            raise

    def update_data(self, collection, old_ids, new_vectors, new_labels=None, new_descs=None):
        """更新数据"""
        try:
            # 先删除旧数据
            self.delete_data(collection, old_ids)

            # 再插入新数据
            self.insert_data(collection, new_vectors, new_labels, new_descs)

            print(f"已更新 {len(old_ids)} 条数据")
            return True
        except Exception as e:
            print(f"更新数据失败: {e}")
            raise

    def search_data(self, collection, query_vec, metric_type, top_k=5):
        """搜索数据"""
        try:
            search_params = {
                "metric_type": metric_type,
                "params": {"nprobe": 10}
            }
            results = collection.search(
                data=[query_vec],
                anns_field="vector",
                param=search_params,
                limit=top_k,
                output_fields=["label", "desc"]
            )
            print(f"\n===== {metric_type} 搜索结果 =====")
            for hits in results:
                for idx, hit in enumerate(hits, 1):
                    print(f"排名 {idx}: "
                          f"距离={hit.distance:.6f}, "
                          f"标签={hit.entity['label']}, "
                          f"描述={hit.entity['desc']}")
            return results
        except Exception as e:
            print(f"搜索失败: {e}")
            raise

    def get_collection_info(self, collection_name):
        """获取集合信息"""
        try:
            collection = Collection(collection_name, using=milvus_alias)
            info = {
                "name": collection.name,
                "num_entities": collection.num_entities,
                "schema": collection.schema.fields,
                "index": collection.indexes
            }
            return info
        except Exception as e:
            print(f"获取集合信息失败: {e}")
            raise


    def close_connection(self):
        """关闭连接"""
        try:
            connections.disconnect(milvus_alias)
            print("已断开Milvus连接")
        except Exception as e:
            print(f"断开连接失败: {e}")


# 使用示例
def main():
    """主函数"""
    try:
        manager = MilvusVectorManager()
        collection = manager.create_collection("L2")

        # 生成测试数据
        base_vec = np.random.rand(dim).astype(np.float32)
        vectors = []
        labels = []
        descs = []

        for group_id in range(3):
            for i in range(10):
                vec = base_vec + np.random.normal(0, 0.05, dim).astype(np.float32)
                vectors.append(vec.tolist())
                labels.append(group_id)
                descs.append(f"组{group_id}_样本{i}")
        manager.insert_data(collection, vectors, labels, descs)
        # manager.delete_data(collection, [0, 1, 2])
        #更新数据库的向量
        # manager.update_data(collection, [0, 1, 2], vectors[3:], labels[3:], descs[3:])
        #查询所有的向量
        # 使用第一个向量作为查询向量
        print(manager.get_collection_info(collection.name))
        query_vec = (base_vec + np.random.normal(0, 0.03, dim).astype(np.float32)).tolist()
        manager.search_data(collection, query_vec, "L2")

        info = manager.get_collection_info(collection.name)
        print(f"\n集合信息: {info}")

    except Exception as e:
        print(f"程序执行出错: {e}")
    finally:
        # 安全关闭连接
        if 'manager' in locals():
            manager.close_connection()



if __name__ == "__main__":
    main()
